package com.smile.commons.tree;

import java.util.*;
import java.util.function.*;

public class TreeBuilder {

    @FunctionalInterface
    public interface ChildrenSetter<T> {
        void setChildren(T parent, List<T> children);
    }

    @FunctionalInterface
    public interface ChildAdder<T> {
        void addChild(T parent, T child);
    }

    public static <T> List<T> buildTree(
            List<T> flatList,
            Function<T, ?> idGetter,
            Function<T, ?> parentIdGetter,
            ChildAdder<T> childAdder,  // 改成了单个子节点添加器
            Predicate<T> isRootCondition) {

        Map<Object, T> nodeMap = new HashMap<>();
        List<T> roots = new ArrayList<>();

        // 1. 创建节点映射并初始化子节点列表
        for (T item : flatList) {
            Object id = idGetter.apply(item);
            nodeMap.put(id, item);
        }

        // 2. 构建树结构
        for (T item : flatList) {
            // 处理根节点
            if (isRootCondition.test(item)) {
                roots.add(item);
                continue;
            }

            // 挂载到父节点
            Object parentId = parentIdGetter.apply(item);
            T parent = nodeMap.get(parentId);
            if (parent != null) {
                // 正确用法：添加单个子节点
                childAdder.addChild(parent, item); // 修复的类型匹配
            }
        }

        return roots;
    }
}